"""模型管理器 - 单例模式管理模型加载"""
import logging
from typing import Optional, Dict, Any
from transformers import (
    AutoTokenizer,
    AutoModel,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification
)
import torch

from app.config import settings

logger = logging.getLogger(__name__)


class ModelManager:
    """模型管理器单例类"""
    
    _instance: Optional['ModelManager'] = None
    _initialized: bool = False
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(ModelManager, cls).__new__(cls)
        return cls._instance
    
    def __init__(self):
        if not self._initialized:
            self.tokenizer: Optional[Any] = None
            self.base_model: Optional[Any] = None
            self.classification_model: Optional[Any] = None
            self.ner_model: Optional[Any] = None
            self.device = settings.device
            self._initialized = True
    
    def load_tokenizer(self):
        """加载分词器"""
        if self.tokenizer is None:
            logger.info(f"正在加载分词器: {settings.LEGALBERT_MODEL_NAME}")
            self.tokenizer = AutoTokenizer.from_pretrained(
                settings.LEGALBERT_MODEL_NAME
            )
            logger.info("分词器加载完成")
        return self.tokenizer
    
    def load_base_model(self):
        """加载基础模型（用于相似度计算）"""
        if self.base_model is None:
            logger.info(f"正在加载基础模型: {settings.LEGALBERT_MODEL_NAME}")
            self.base_model = AutoModel.from_pretrained(
                settings.LEGALBERT_MODEL_NAME
            )
            self.base_model.to(self.device)
            self.base_model.eval()
            logger.info(f"基础模型加载完成，设备: {self.device}")
        return self.base_model
    
    def load_classification_model(self):
        """加载分类模型"""
        if self.classification_model is None:
            logger.info(f"正在加载分类模型: {settings.LEGALBERT_MODEL_NAME}")
            # 注意：实际使用时需要根据分类任务微调模型
            # 这里使用基础模型，可以根据需要替换为微调后的模型路径
            try:
                self.classification_model = AutoModelForSequenceClassification.from_pretrained(
                    settings.LEGALBERT_MODEL_NAME,
                    num_labels=2  # 默认二分类，可根据实际需求调整
                )
            except Exception as e:
                logger.warning(f"无法加载分类模型，使用基础模型: {e}")
                # 如果模型不支持分类，使用基础模型
                self.classification_model = AutoModel.from_pretrained(
                    settings.LEGALBERT_MODEL_NAME
                )
            self.classification_model.to(self.device)
            self.classification_model.eval()
            logger.info(f"分类模型加载完成，设备: {self.device}")
        return self.classification_model
    
    def load_ner_model(self):
        """加载 NER 模型"""
        if self.ner_model is None:
            logger.info(f"正在加载 NER 模型: {settings.LEGALBERT_MODEL_NAME}")
            # 注意：实际使用时需要根据 NER 任务微调模型
            # 这里使用基础模型，可以根据需要替换为微调后的模型路径
            try:
                self.ner_model = AutoModelForTokenClassification.from_pretrained(
                    settings.LEGALBERT_MODEL_NAME,
                    num_labels=9  # 默认标签数，可根据实际需求调整（BIO 标注）
                )
            except Exception as e:
                logger.warning(f"无法加载 NER 模型，使用基础模型: {e}")
                # 如果模型不支持 NER，使用基础模型
                self.ner_model = AutoModel.from_pretrained(
                    settings.LEGALBERT_MODEL_NAME
                )
            self.ner_model.to(self.device)
            self.ner_model.eval()
            logger.info(f"NER 模型加载完成，设备: {self.device}")
        return self.ner_model
    
    def initialize_models(self):
        """初始化所有模型"""
        logger.info("开始初始化所有模型...")
        self.load_tokenizer()
        self.load_base_model()
        self.load_classification_model()
        self.load_ner_model()
        logger.info("所有模型初始化完成")


# 全局模型管理器实例
model_manager = ModelManager()

